import pygame
import os
import threading
import json
import socket
import ipaddress

UDP_MSG = "Wait For MSG"
# 设置颜色
RED = (255, 0, 0)
GREEN = (0, 255, 0)
# 需要控制的管脚代号
pins = [4,5,6,7,0,1,2,3]

INTERVAL_X = 100
INTERVAL_Y = 100
columns_num = 4

# 设置圆球的半径和位置
radius = 50
positions = []
for k, p in enumerate(pins):
    positions.append((INTERVAL_X * (k % columns_num + 1), INTERVAL_Y * (k // columns_num + 1)))

# 设置8个圆球的状态，默认为红色
CIRCLES = [RED] * 8

def is_valid_ip(ip_str):
    try:
        # 使用ipaddress库的ip_address方法尝试解析IP地址
        ipaddress.ip_address(ip_str)
        return True
    except ValueError:
        return False

if os.path.exists("config.json"):
    with open("config.json", "r") as f:
        config:dict = json.load(f)
    if "LastRemoteIP" not in config.keys():
        config["LastRemoteIP"] = "0.0.0.0"
else:
    config = {}
    config["LastRemoteIP"] = "0.0.0.0"

# 本地IP和端口
# 获取本地主机名
hostname = socket.gethostname()
# 获取本地IP地址
localIP = socket.gethostbyname(hostname)
localPort = 8888
print(f"localIP:{localIP}, localPort: {localPort}")
print(f"LastIP:{config['LastRemoteIP']}")
# 监听UDP数据包的线程函数
def listen_thread():
    global UDP_MSG, CIRCLES
    states = [RED, GREEN]
    while True:
        # 从UDP套接字接收数据
        data, addr = sock.recvfrom(1024)
        try:
            _data = data.decode()
        except Exception as e:
            _data = str(e)
        UDP_MSG = "from {} :{}".format(addr[0], _data)
        print("\nfrom {}:\n{}\n".format(addr[0], _data), end="")
        if _data.startswith("Pin"):
            _, pin, _, state, _ = _data.split()
            CIRCLES[pins.index(int(pin))] = states[int(state)]


sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.bind((localIP, localPort))
threading.Thread(target=listen_thread).start()

# 远程IP和端口（ESP32的IP和端口）
while True:
    print(f"直接回车可填入上次连接的地址")
    remoteIP = input("输入机器人的IP地址: ")
    if is_valid_ip(remoteIP):
        print(f"\"{remoteIP}\" 是一个有效的IP地址")
        config["LastRemoteIP"] = remoteIP
        break
    elif remoteIP:
        print(f"\"{remoteIP}\" 不是一个有效的IP地址")
    else:
        remoteIP = config["LastRemoteIP"]
        if is_valid_ip(remoteIP):
            print(f"\"{remoteIP}\" 是一个有效的IP地址")
        break

remotePort = 8888
# 创建UDP套接字
sock.sendto(f"$0 {localIP}".encode(), (remoteIP, remotePort))



# 初始化Pygame
pygame.init()
# 设置窗口大小
window_size = (positions[columns_num-1][0]+100, positions[-1][-1]+100)
screen = pygame.display.set_mode(window_size)
# 设置字体和字号
font = pygame.font.SysFont(None, 30)





while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            with open("config.json", "w") as f:
                json.dump(config, f, indent=4)
            pygame.quit()
            os._exit(0)
        elif event.type == pygame.MOUSEBUTTONDOWN:
            # 获取鼠标点击位置
            mouse_x, mouse_y = pygame.mouse.get_pos()
            # 检查是否点击到了圆球
            for i, (x, y) in enumerate(positions):
                if (mouse_x - x)**2 + (mouse_y - y)**2 < radius**2:
                    # 点击后切换圆球颜色
                    if CIRCLES[i] == RED:
                        # CIRCLES[i] = green
                        sock.sendto(f":{pins[i]} 1".encode(), (remoteIP, remotePort))
                    else:
                        # CIRCLES[i] = RED
                        sock.sendto(f":{pins[i]} 0".encode(), (remoteIP, remotePort))

    # 清空屏幕
    screen.fill((255, 255, 255))

    # 绘制8个圆球和其中的数字
    for i, (x, y) in enumerate(positions):
        pygame.draw.circle(screen, CIRCLES[i], (x, y), radius)
        text_surface = font.render(str(pins[i]), True, (0, 0, 0))
        text_rect = text_surface.get_rect(center=(x, y))
        screen.blit(text_surface, text_rect)
    
    text_surface = font.render(UDP_MSG[:-1], True, (0, 0, 0))
    text_rect = text_surface.get_rect(center=(positions[0][0]+100, positions[-1][-1]+80))
    
    screen.blit(text_surface, text_rect)
    # 更新显示
    pygame.display.flip()